#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据库迁移脚本 - 添加班级和作业功能
"""

import sqlite3
import os
from datetime import datetime

def migrate_database():
    """执行数据库迁移"""
    db_path = 'oj_system.db'
    
    if not os.path.exists(db_path):
        print("数据库文件不存在，请先运行主程序创建数据库")
        return
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("开始数据库迁移...")
        
        # 检查并创建班级表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS classes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name VARCHAR(100) NOT NULL,
                description TEXT,
                teacher_id INTEGER NOT NULL,
                max_students INTEGER DEFAULT 50,
                status VARCHAR(20) DEFAULT 'active',
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (teacher_id) REFERENCES users (id)
            )
        """)
        print("✓ 班级表创建完成")
        
        # 检查并创建班级-学生关联表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS class_student (
                class_id INTEGER NOT NULL,
                student_id INTEGER NOT NULL,
                PRIMARY KEY (class_id, student_id),
                FOREIGN KEY (class_id) REFERENCES classes (id),
                FOREIGN KEY (student_id) REFERENCES users (id)
            )
        """)
        print("✓ 班级-学生关联表创建完成")
        
        # 检查并创建作业表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS assignments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title VARCHAR(200) NOT NULL,
                description TEXT,
                class_id INTEGER NOT NULL,
                problem_ids TEXT,
                due_date DATETIME NOT NULL,
                total_score FLOAT DEFAULT 100.0,
                status VARCHAR(20) DEFAULT 'active',
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (class_id) REFERENCES classes (id)
            )
        """)
        print("✓ 作业表创建完成")
        
        # 检查并创建作业提交表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS assignment_submissions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                assignment_id INTEGER NOT NULL,
                student_id INTEGER NOT NULL,
                problem_id INTEGER NOT NULL,
                code TEXT,
                language VARCHAR(20) DEFAULT 'python',
                status VARCHAR(20) DEFAULT 'pending',
                score FLOAT DEFAULT 0.0,
                execution_time FLOAT,
                memory_usage FLOAT,
                error_message TEXT,
                submitted_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                graded_at DATETIME,
                graded_by INTEGER,
                feedback TEXT,
                FOREIGN KEY (assignment_id) REFERENCES assignments (id),
                FOREIGN KEY (student_id) REFERENCES users (id),
                FOREIGN KEY (problem_id) REFERENCES problems (id),
                FOREIGN KEY (graded_by) REFERENCES users (id)
            )
        """)
        print("✓ 作业提交表创建完成")
        
        # 创建索引
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_classes_teacher ON classes (teacher_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_classes_status ON classes (status)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_assignments_class ON assignments (class_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_assignments_status ON assignments (status)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_assignment_submissions_assignment ON assignment_submissions (assignment_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_assignment_submissions_student ON assignment_submissions (student_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_assignment_submissions_problem ON assignment_submissions (problem_id)")
        print("✓ 索引创建完成")
        
        # 提交事务
        conn.commit()
        print("✓ 数据库迁移完成！")
        
        # 显示表信息
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name LIKE '%class%' OR name LIKE '%assignment%'")
        tables = cursor.fetchall()
        print(f"\n创建的表: {[table[0] for table in tables]}")
        
    except Exception as e:
        conn.rollback()
        print(f"❌ 迁移失败: {e}")
        raise
    finally:
        conn.close()

def create_sample_data():
    """创建示例数据"""
    db_path = 'oj_system.db'
    
    if not os.path.exists(db_path):
        print("数据库文件不存在")
        return
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("\n开始创建示例数据...")
        
        # 检查是否有教师用户
        cursor.execute("SELECT id, name FROM users WHERE role = 'teacher' LIMIT 1")
        teacher = cursor.fetchone()
        
        if not teacher:
            print("⚠️  没有找到教师用户，请先创建教师账户")
            return
        
        teacher_id = teacher[0]
        teacher_name = teacher[1]
        
        # 创建示例班级
        cursor.execute("""
            INSERT OR IGNORE INTO classes (name, description, teacher_id, max_students, status)
            VALUES (?, ?, ?, ?, ?)
        """, (f"{teacher_name}的编程基础班", "学习Python编程基础", teacher_id, 30, 'active'))
        
        class_id = cursor.lastrowid
        print(f"✓ 创建示例班级: {teacher_name}的编程基础班")
        
        # 获取学生用户
        cursor.execute("SELECT id, name FROM users WHERE role = 'student' LIMIT 5")
        students = cursor.fetchall()
        
        if students:
            # 添加学生到班级
            for student_id, student_name in students:
                cursor.execute("""
                    INSERT OR IGNORE INTO class_student (class_id, student_id)
                    VALUES (?, ?)
                """, (class_id, student_id))
            print(f"✓ 添加 {len(students)} 名学生到班级")
        
        # 检查是否有题目
        cursor.execute("SELECT id, title FROM problems LIMIT 3")
        problems = cursor.fetchall()
        
        if problems:
            problem_ids = [str(problem[0]) for problem in problems]
            
            # 创建示例作业
            due_date = datetime.now().replace(hour=23, minute=59, second=59)
            cursor.execute("""
                INSERT OR IGNORE INTO assignments (title, description, class_id, problem_ids, due_date, total_score, status)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                "Python基础练习",
                "完成以下Python编程题目，巩固基础知识",
                class_id,
                ','.join(problem_ids),
                due_date,
                100.0,
                'active'
            ))
            print("✓ 创建示例作业: Python基础练习")
        
        conn.commit()
        print("✓ 示例数据创建完成！")
        
    except Exception as e:
        conn.rollback()
        print(f"❌ 创建示例数据失败: {e}")
        raise
    finally:
        conn.close()

if __name__ == "__main__":
    print("=== 班级和作业功能数据库迁移 ===")
    
    try:
        migrate_database()
        create_sample_data()
        print("\n🎉 迁移完成！现在可以启动后端服务了。")
        
    except Exception as e:
        print(f"\n❌ 迁移失败: {e}")
        import traceback
        traceback.print_exc()
